# -*-coding:utf-8-*-
import glob
import os.path
import numpy as np
import tensorflow as tf
from tensorflow.python.platform import gfile
import tensorflow.contrib.slim as slim
#y也可以在代码中用slim=tf.contrib.slim来代替上面加载slim库

#加载通过Tensorflow-Slim定义好的inception-v3模型
import tensorflow.contrib.slim.python.slim.nets.inception_v3 as inception_v3

#处理好后的数据文件
input_data='path/flower_processed_data.npy'   #make_data.py中做好的
#保存训练好的模型的路径。这里可以将使用新数据训练得到的完整模型保存下来，
# 如果计算资源充足，还可以在训练完最后的全连接层之后再训练所有的网络层，这样可以使得模型更加贴近新数据
train_file='path/save_model'
#谷歌提供的训练好的模型文件地址
ckpt_file='path/inception_v3.ckpt'

#定义训练中使用的参数
learning_rate=0.0001
steps=300
batch=32
n_classes=5

#不需要从谷歌训练好的模型中加载的参数，这里指最后的全连接层，因为新的问题中要重新训练这一层中的参数，这里给出的是参数的前缀
#不需要加载的就是要重新训练的
checkpoint_exclude_scopes='InceptionV3/Logits,InceptionV3/Auxlogits'
trainable_scopes='InceptionV3/Logits,InceptionV3/AuxLogits'

#获取所有需要从谷歌训练好的模型中加载的参数
def get_tuned_variables():
    #排除列表
    exclusions=[scope.strip() for scope in checkpoint_exclude_scopes.split(',')]

    variables_to_restore=[]
    #枚举inception-v3中所有的参数，然后判断是否需要从加载列表中移除
    #slim.get_model_variables()是得到变量的列表
    for var in slim.get_model_variables():
        excluded=False
        for exclusion in exclusions:
            #检查字符串是否以以指定字符串开头,最终将不以其开头的变量加入到集合中
            if var.op.name.startswith(exclusion):
                excluded=True
                break
        if not excluded:
            variables_to_restore.append(var)
    #返回需要加载的参数
    return variables_to_restore

#获取所有需要训练的变量列表,不需要加载的参数就是需要重新训练的参数
def get_trainable_variables():
    scopes=[scope.strip() for scope in trainable_scopes.split(',')]
    variables_to_train=[]
    #枚举所有需要训练的参数前缀，并通过这些前缀找到所有的参数
    for scope in scopes:
        variables=tf.get_collection(
            tf.GraphKeys.TRAINABLE_VARIABLES,scope
        )
        variables_to_train.extend(variables)     #列表加元素（列表的合并）
    return variables_to_train

def main():
    #加载预处理好的数据
    processed_data=np.load(input_data)
    training_images=processed_data[0]
    training_labels=processed_data[1]
    validation_images=processed_data[3]
    validation_labels=processed_data[4]
    testing_images=processed_data[5]
    testing_labels=processed_data[6]

    n_training_example=len(training_images)
    print("%d training example,%d validation example and %d testing example" \
          %(n_training_example,len(validation_labels),len(testing_labels)))

    #定义inception-v3的输入，images为输入图片，labels为每一张图片对应的标签。
    images=tf.placeholder(tf.float32,[None,299,299,3],name='input_images')
    labels=tf.placeholder(tf.int64,[None],name='labels')

    #定义inception-v3模型。因为谷歌给出的只有模型参数取值，所以这里需要在这个代码中定义inception-v3的模型结构。
    # 虽然理论上需要区分训练与测试使用的模型，如测试时需要提供参数is_training=False，
    # 但是因为预先训练好的inception-v3模型中使用的batch normalization参数与新的数据会有差异，导致结果很差，
    # 所以这里直接使用同一个模型来进行测试
    with slim.arg_scope(inception_v3.inception_v3_arg_scope()):
        logits,_=inception_v3.inception_v3(images,num_classes=n_classes)

    #获取需要训练的变量
    trainable_variables=get_trainable_variables()
    #定义交叉熵损失，注意在模型定义的时候已经将正则化损失加入损失集合了
    tf.losses.softmax_cross_entropy(
        tf.one_hot(labels,n_classes),logits,weights=1.0
    )
    #定义训练过程，这里minimize的过程中指定了需要优化的变量集合
    train_step=tf.train.RMSPropOptimizer(learning_rate).minimize(tf.losses.get_total_loss)
    #计算正确率
    with tf.name_scope('evaluation'):
        correct_prediction=tf.equal(tf.argmax(logits,1),labels)
        evaluation_step=tf.reduce_mean(tf.cast(correct_prediction,tf.float32))

    #定义加载模型的函数
    load_fn=slim.assign_from_checkpoint_fn(
        ckpt_file,
        get_tuned_variables(),
        ignore_missing_vars=True
    )

    #定义保存新的训练好的模型的函数
    saver=tf.train.Saver()
    with tf.Session() as sess:
        #初始化没有加载进来的变量。注意这个过程一定要在模型加载之前，否则初始化过程会将已经加载好的变量重新赋值
        init=tf.global_variables_initializer()
        sess.run(init)

        #加载谷歌已经训练好的模型
        print('Loading tuned variables from %s'%(ckpt_file))
        load_fn(sess)

        start=0
        end=batch
        for i in range(steps):
            #运行训练过程，这里不会更新全部的参数，只会更新部分指定的参数
            sess.run(train_step,feed_dict={
                images:training_images[start:end],
                labels:training_labels[start:end]
            })

            #输出日志
            if i%30==0 or i+1==steps:
                saver.save(sess,train_file,global_step=i)
                validation_accuracy=sess.run(evaluation_step,feed_dict={
                    images:validation_images,labels:validation_labels
                })
                print('Step %d:Validation accuracy=%.1f%%'%(i,validation_accuracy*100.0))

            #因为在数据预处理的时候已经做过了打乱数据的操作，所以这里只需要顺序使用训练数据就好
            start=end
            if start==n_training_example:
                start=0

            end=start+batch
            if end>n_training_example:
                end=n_training_example

        #在最后的测试数据上测试正确率
        test_accuracy=sess.run(evaluation_step,feed_dict={
            images:testing_images,labels:testing_labels
        })
        print('Final test accuracy=%.1f%%'%(test_accuracy*100.0))

if __name__=='__main__':
    tf.app.run()